{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import re\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "import art\n",
    "from art.local import LocalBackend\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "BASE_MODEL = \"Qwen/Qwen2.5-7B-Instruct\"\n",
    "PRISONERS_DILEMMA_ROUNDS = 10\n",
    "TRAINING_STEPS = 1_000\n",
    "\n",
    "backend = LocalBackend()\n",
    "model = art.TrainableModel(\n",
    "    name=\"001\", project=\"prisoners-dilemma\", base_model=BASE_MODEL\n",
    ")\n",
    "await model.register(backend)\n",
    "\n",
    "client = model.openai_client()\n",
    "\n",
    "\n",
    "prompt = f\"\"\"\n",
    "You are playing a game of prisoners' dilemma with another player.\n",
    "\n",
    "You are given a choice between two actions:\n",
    "\n",
    "1. Cooperate\n",
    "2. Defect\n",
    "\n",
    "The payoffs are as follows:\n",
    "\n",
    "- If both players cooperate, you get 3 points and the other player gets 3 points.\n",
    "- If one player cooperates and the other defects, the defector gets 5 points and the cooperator gets 0 points.\n",
    "- If both players defect, you both get 1 point.\n",
    "\n",
    "You will play this game {PRISONERS_DILEMMA_ROUNDS} times with the same player.\n",
    "\n",
    "For your first turn, would you like to cooperate or defect?\".\n",
    "\"\"\".strip()\n",
    "\n",
    "\n",
    "async def rollout_game(\n",
    "    models: tuple[str, str] = (model.name, model.name),\n",
    ") -> tuple[art.Trajectory, art.Trajectory]:\n",
    "    messages: tuple[art.Messages, art.Messages] = (\n",
    "        [{\"role\": \"user\", \"content\": prompt}],\n",
    "        [{\"role\": \"user\", \"content\": prompt}],\n",
    "    )\n",
    "    trajectories = (\n",
    "        art.Trajectory(messages_and_choices=[*messages[0]], reward=0),\n",
    "        art.Trajectory(messages_and_choices=[*messages[1]], reward=0),\n",
    "    )\n",
    "    for _ in range(PRISONERS_DILEMMA_ROUNDS):\n",
    "        chat_completions = await asyncio.gather(\n",
    "            client.chat.completions.create(\n",
    "                messages=messages[0], model=models[0], max_completion_tokens=512\n",
    "            ),\n",
    "            client.chat.completions.create(\n",
    "                messages=messages[1], model=models[1], max_completion_tokens=512\n",
    "            ),\n",
    "        )\n",
    "        choices = [chat_completion.choices[0] for chat_completion in chat_completions]\n",
    "        messages[0].append({\"role\": \"assistant\", \"content\": choices[0].message.content})\n",
    "        messages[1].append({\"role\": \"assistant\", \"content\": choices[1].message.content})\n",
    "        trajectories[0].messages_and_choices.append(choices[0])\n",
    "        trajectories[1].messages_and_choices.append(choices[1])\n",
    "        actions = [\n",
    "            (\n",
    "                matches[-1]\n",
    "                if (\n",
    "                    matches := re.findall(\n",
    "                        pattern=r\"cooperate|defect\",\n",
    "                        string=(choice.message.content or \"\").lower(),\n",
    "                    )\n",
    "                )\n",
    "                else \"none\"\n",
    "            )\n",
    "            for choice in choices\n",
    "        ]\n",
    "        if actions[0] == \"cooperate\" and actions[1] == \"cooperate\":\n",
    "            trajectories[0].reward += 3\n",
    "            trajectories[1].reward += 3\n",
    "        elif actions[0] == \"cooperate\" and actions[1] == \"defect\":\n",
    "            trajectories[0].reward += 0\n",
    "            trajectories[1].reward += 5\n",
    "        elif actions[0] == \"defect\" and actions[1] == \"cooperate\":\n",
    "            trajectories[0].reward += 5\n",
    "            trajectories[1].reward += 0\n",
    "        elif actions[0] == \"defect\" and actions[1] == \"defect\":\n",
    "            trajectories[0].reward += 1\n",
    "            trajectories[1].reward += 1\n",
    "        else:\n",
    "            # One or both players did not choose an action.\n",
    "            default_rewards = {\"cooperate\": 3, \"defect\": 5, \"none\": 0}\n",
    "            trajectories[0].reward += default_rewards[actions[0]]\n",
    "            trajectories[1].reward += default_rewards[actions[1]]\n",
    "        for i in range(2):\n",
    "            joiner = \"\\n> \"\n",
    "            messages[i].append(\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": f\"The other player responded as follows: \\n\\n> {joiner.join((choices[1 - i].message.content or '').splitlines())}\\n\\n\"\n",
    "                    f\"Your score is {trajectories[i].reward}. The other player's score is {trajectories[1 - i].reward}.\\n\\n\"\n",
    "                    \"For the next round, would you like to cooperate or defect?\",\n",
    "                }\n",
    "            )\n",
    "            trajectories[i].messages_and_choices.append(messages[i][-1])\n",
    "    return trajectories\n",
    "\n",
    "\n",
    "for _ in range(await model.get_step(), TRAINING_STEPS):\n",
    "    # Simultaneously rollout self-play games, and games versus the base model.\n",
    "    self_play_trajectories, base_play_trajectories = await asyncio.gather(\n",
    "        art.gather_trajectories(\n",
    "            (rollout_game(models=(model.name, model.name)) for _ in range(8)),\n",
    "            pbar_desc=\"versus-self\",\n",
    "        ),\n",
    "        art.gather_trajectories(\n",
    "            (rollout_game(models=(model.name, BASE_MODEL)) for _ in range(8)),\n",
    "            pbar_desc=\"versus-base\",\n",
    "        ),\n",
    "    )\n",
    "    # Log performance versus self and the base model, as well as the base model's performance.\n",
    "    await model.log(\n",
    "        [t for ts in self_play_trajectories for t in ts], split=\"versus-self\"\n",
    "    )\n",
    "    await model.log([ts[0] for ts in base_play_trajectories], split=\"versus-base\")\n",
    "    await model.log([ts[1] for ts in base_play_trajectories], split=\"base-model\")\n",
    "    # Train the model on self-play and base-play trajectories.\n",
    "    await model.train(\n",
    "        trajectory_groups=[\n",
    "            # Since all self-play games have the same starting state and are symmetric, we can gather\n",
    "            # trajectories from all self-play games into a single trajectory group.\n",
    "            art.TrajectoryGroup(t for ts in self_play_trajectories for t in ts),\n",
    "            # We can also gather all base-play _trained model_ trajectories into a single trajectory group.\n",
    "            # We don't want to train on base model trajectories, because they are sampled from a different distribution.\n",
    "            art.TrajectoryGroup(ts[0] for ts in base_play_trajectories),\n",
    "        ],\n",
    "        config=art.TrainConfig(learning_rate=5e-5),\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
